#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
科学绘图代码 - SCI期刊风格
功能：
1. 小提琴图（多组对比）
2. 蜂窝图（子图）
3. 散点气泡图
4. 发散型条形图
5. 连续变量直方图（多组对比）
6. 二维渐变色地形图
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.font_manager import FontProperties
import matplotlib.patches as patches
from matplotlib.colors import LinearSegmentedColormap
import os
import warnings
warnings.filterwarnings('ignore')

class ScientificPlotter:
    def __init__(self):
        """初始化科学绘图器"""
        self.setup_chinese_font()
        self.setup_color_schemes()
        self.output_dir = "数据预处理_绘图结果"
        
        # 确保输出目录存在
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)
    
    def setup_chinese_font(self):
        """设置中文字体"""
        font_paths = [
            'C:/Windows/Fonts/simhei.ttf',  # 黑体（优先）
            'C:/Windows/Fonts/msyh.ttc',   # 微软雅黑
            'C:/Windows/Fonts/simsun.ttc', # 宋体
        ]
        
        self.font_prop = None
        for path in font_paths:
            if os.path.exists(path):
                try:
                    self.font_prop = FontProperties(fname=path)
                    plt.rcParams['font.family'] = self.font_prop.get_name()
                    plt.rcParams['axes.unicode_minus'] = False  # 修复负号显示
                    print(f"成功加载字体: {path}")
                    break
                except Exception as e:
                    print(f"加载字体失败 {path}: {e}")
                    continue
        
        if self.font_prop is None:
            print("警告：未找到可用的中文字体，可能显示异常")
            self.font_prop = FontProperties()
    
    def setup_color_schemes(self):
        """设置配色方案"""
        # 第一种渐变色系（绿色系）
        self.colors_green = [
            (0/255, 70/255, 41/255),
            (12/255, 113/255, 59/255),
            (55/255, 158/255, 84/255),
            (119/255, 197/255, 120/255),
            (186/255, 226/255, 148/255),
            (236/255, 247/255, 177/255),
            (254/255, 254/255, 227/255)
        ]
        
        # 第二种渐变色系（蓝色系）
        self.colors_blue = [
            (10/255, 31/255, 94/255),
            (34/255, 65/255, 153/255),
            (29/255, 128/255, 185/255),
            (62/255, 179/255, 196/255),
            (144/255, 212/255, 185/255),
            (218/255, 240/255, 178/255),
            (252/255, 253/255, 211/255)
        ]
        
        # 第三种渐变色系（紫黄系）
        self.colors_purple = [
            (78/255, 98/255, 171/255),
            (70/255, 158/255, 180/255),
            (135/255, 207/255, 164/255),
            (203/255, 233/255, 137/255),
            (245/255, 251/255, 177/255),
            (254/255, 254/255, 154/255),
            (253/255, 185/255, 106/255)
        ]
        
        # 创建自定义颜色映射
        self.cmap_green = LinearSegmentedColormap.from_list("custom_green", self.colors_green)
        self.cmap_blue = LinearSegmentedColormap.from_list("custom_blue", self.colors_blue)
        self.cmap_purple = LinearSegmentedColormap.from_list("custom_purple", self.colors_purple)
    
    def set_plot_style(self):
        """设置绘图样式"""
        plt.rcParams.update({
            'figure.dpi': 300,
            'savefig.dpi': 300,
            'font.size': 10,
            'axes.linewidth': 1.2,
            'axes.spines.top': False,
            'axes.spines.right': False,
            'axes.grid': True,
            'grid.alpha': 0.3,
            'grid.linestyle': '--',
            'grid.linewidth': 0.5
        })
    
    def load_processed_data(self):
        """加载处理后的数据"""
        try:
            self.male_data = pd.read_csv('男胎检测数据_处理后.csv')
            self.female_data = pd.read_csv('女胎检测数据_处理后.csv')
            
            # 处理检测孕周列，将"11w+6"格式转换为数值
            def convert_pregnancy_week(week_str):
                try:
                    if pd.isna(week_str) or week_str == '':
                        return np.nan
                    week_str = str(week_str).strip()
                    if 'w' in week_str:
                        parts = week_str.split('w')
                        weeks = float(parts[0])
                        days = 0
                        if len(parts) > 1 and '+' in parts[1]:
                            days = float(parts[1].replace('+', ''))
                        return weeks + days / 7.0
                    else:
                        return float(week_str)
                except:
                    return np.nan
            
            # 转换检测孕周列
            if '检测孕周' in self.male_data.columns:
                self.male_data['检测孕周'] = self.male_data['检测孕周'].apply(convert_pregnancy_week)
            if '检测孕周' in self.female_data.columns:
                self.female_data['检测孕周'] = self.female_data['检测孕周'].apply(convert_pregnancy_week)
            
            # 确保数值列为数值类型
            numeric_columns = ['年龄', '身高', '体重', '孕妇BMI', '检测孕周', '原始读段数', 
                              '在参考基因组上比对的比例', '重复读段的比例', 'GC含量', 
                              '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值', 
                              'X染色体的Z值', 'Y染色体的Z值', 'Y染色体浓度', 'X染色体浓度',
                              '13号染色体的GC含量', '18号染色体的GC含量', '21号染色体的GC含量',
                              '被过滤掉读段数的比例', '怀孕次数', '生产次数']
            
            for col in numeric_columns:
                if col in self.male_data.columns:
                    self.male_data[col] = pd.to_numeric(self.male_data[col], errors='coerce')
                if col in self.female_data.columns:
                    self.female_data[col] = pd.to_numeric(self.female_data[col], errors='coerce')
            
            print(f"成功加载数据: 男胎 {self.male_data.shape}, 女胎 {self.female_data.shape}")
            return True
        except Exception as e:
            print(f"数据加载失败: {e}")
            return False
    
    def plot_violin_comparison(self):
        """绘制小提琴图（多组对比）"""
        print("正在绘制小提琴图...")
        self.set_plot_style()
        
        # 准备数据
        variables = ['年龄', '身高', '体重', '孕妇BMI', '检测孕周']
        
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        axes = axes.flatten()
        
        for i, var in enumerate(variables):
            if var in self.male_data.columns and var in self.female_data.columns:
                # 准备绘图数据
                data_plot = []
                labels = []
                
                # 男胎数据
                male_vals = self.male_data[var].dropna()
                data_plot.extend(male_vals.tolist())
                labels.extend(['男胎'] * len(male_vals))
                
                # 女胎数据
                female_vals = self.female_data[var].dropna()
                data_plot.extend(female_vals.tolist())
                labels.extend(['女胎'] * len(female_vals))
                
                # 创建DataFrame用于绘图
                plot_df = pd.DataFrame({'值': data_plot, '性别': labels})
                
                # 绘制小提琴图
                violin_parts = axes[i].violinplot([male_vals, female_vals], 
                                                 positions=[0, 1], 
                                                 showmeans=True, 
                                                 showmedians=True)
                
                # 设置颜色
                colors = [self.colors_green[2], self.colors_blue[2]]
                for pc, color in zip(violin_parts['bodies'], colors):
                    pc.set_facecolor(color)
                    pc.set_alpha(0.7)
                
                # 设置标签和标题
                axes[i].set_xticks([0, 1])
                axes[i].set_xticklabels(['男胎', '女胎'], fontproperties=self.font_prop)
                axes[i].set_ylabel(f'{var}', fontproperties=self.font_prop)
                axes[i].set_title(f'{var}分布对比', fontproperties=self.font_prop, fontsize=12)
                axes[i].grid(True, alpha=0.3)
        
        # 隐藏多余的子图
        axes[-1].set_visible(False)
        
        plt.tight_layout()
        plt.savefig(f'{self.output_dir}/小提琴图_多组对比.png', bbox_inches='tight', dpi=300)
        plt.close()
        
        # 分别生成子图
        for i, var in enumerate(variables):
            if var in self.male_data.columns and var in self.female_data.columns:
                fig, ax = plt.subplots(figsize=(6, 4))
                
                male_vals = self.male_data[var].dropna()
                female_vals = self.female_data[var].dropna()
                
                violin_parts = ax.violinplot([male_vals, female_vals], 
                                           positions=[0, 1], 
                                           showmeans=True, 
                                           showmedians=True)
                
                colors = [self.colors_green[2], self.colors_blue[2]]
                for pc, color in zip(violin_parts['bodies'], colors):
                    pc.set_facecolor(color)
                    pc.set_alpha(0.7)
                
                ax.set_xticks([0, 1])
                ax.set_xticklabels(['男胎', '女胎'], fontproperties=self.font_prop)
                ax.set_ylabel(f'{var}', fontproperties=self.font_prop)
                ax.set_title(f'{var}分布对比', fontproperties=self.font_prop, fontsize=12)
                ax.grid(True, alpha=0.3)
                
                plt.tight_layout()
                plt.savefig(f'{self.output_dir}/图{i+1}a_{var}_小提琴图.png', bbox_inches='tight', dpi=300)
                plt.close()
        
        print("小提琴图绘制完成！")
    
    def plot_hexbin_subplots(self):
        """绘制蜂窝图（子图）"""
        print("正在绘制蜂窝图...")
        self.set_plot_style()
        
        # 选择变量对
        var_pairs = [
            ('年龄', '孕妇BMI'),
            ('身高', '体重'),
            ('13号染色体的Z值', '18号染色体的Z值'),
            ('21号染色体的Z值', 'X染色体的Z值')
        ]
        
        # 为男胎和女胎分别绘制蜂窝图
        for data, data_name, color_scheme in [(self.male_data, '男胎', self.cmap_green), 
                                              (self.female_data, '女胎', self.cmap_blue)]:
            
            fig, axes = plt.subplots(2, 2, figsize=(12, 10))
            axes = axes.flatten()
            
            for i, (var1, var2) in enumerate(var_pairs):
                if var1 in data.columns and var2 in data.columns:
                    # 获取两列都非空的数据
                    valid_data = data[[var1, var2]].dropna()
                    if len(valid_data) > 0:
                        x_data = valid_data[var1]
                        y_data = valid_data[var2]
                    else:
                        continue
                    
                    # 绘制蜂窝图
                    hb = axes[i].hexbin(x_data, y_data, gridsize=20, cmap=color_scheme, alpha=0.8)
                    
                    axes[i].set_xlabel(f'{var1}', fontproperties=self.font_prop)
                    axes[i].set_ylabel(f'{var2}', fontproperties=self.font_prop)
                    axes[i].set_title(f'{data_name} - {var1} vs {var2}', fontproperties=self.font_prop)
                    
                    # 添加颜色条
                    cb = plt.colorbar(hb, ax=axes[i])
                    cb.set_label('密度', fontproperties=self.font_prop)
            
            plt.tight_layout()
            plt.savefig(f'{self.output_dir}/蜂窝图_{data_name}_组图.png', bbox_inches='tight', dpi=300)
            plt.close()
            
            # 分别生成子图
            for i, (var1, var2) in enumerate(var_pairs):
                if var1 in data.columns and var2 in data.columns:
                    fig, ax = plt.subplots(figsize=(6, 5))
                    
                    # 获取两列都非空的数据
                    valid_data = data[[var1, var2]].dropna()
                    if len(valid_data) == 0:
                        plt.close()
                        continue
                    x_data = valid_data[var1]
                    y_data = valid_data[var2]
                    
                    hb = ax.hexbin(x_data, y_data, gridsize=20, cmap=color_scheme, alpha=0.8)
                    
                    ax.set_xlabel(f'{var1}', fontproperties=self.font_prop)
                    ax.set_ylabel(f'{var2}', fontproperties=self.font_prop)
                    ax.set_title(f'{data_name} - {var1} vs {var2}', fontproperties=self.font_prop)
                    
                    cb = plt.colorbar(hb, ax=ax)
                    cb.set_label('密度', fontproperties=self.font_prop)
                    
                    plt.tight_layout()
                    plt.savefig(f'{self.output_dir}/图{i+1}{["a", "b"][0 if data_name=="男胎" else 1]}_{data_name}_{var1}_{var2}_蜂窝图.png', 
                               bbox_inches='tight', dpi=300)
                    plt.close()
        
        print("蜂窝图绘制完成！")
    
    def plot_scatter_bubble(self):
        """绘制散点气泡图"""
        print("正在绘制散点气泡图...")
        self.set_plot_style()
        
        # 选择三个维度的变量
        x_var = '年龄'
        y_var = '孕妇BMI'
        size_var = '检测孕周'
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
        
        datasets = [(self.male_data, '男胎', self.colors_green), 
                   (self.female_data, '女胎', self.colors_blue)]
        
        for i, (data, name, colors) in enumerate(datasets):
            if all(var in data.columns for var in [x_var, y_var, size_var]):
                # 获取三列都非空的数据
                valid_data = data[[x_var, y_var, size_var]].dropna()
                if len(valid_data) == 0:
                    continue
                x_data = valid_data[x_var]
                y_data = valid_data[y_var]
                size_data = valid_data[size_var]
                
                # 标准化气泡大小
                size_normalized = (size_data - size_data.min()) / (size_data.max() - size_data.min())
                bubble_sizes = 20 + size_normalized * 200  # 气泡大小范围 20-220
                
                # 绘制散点图
                scatter = axes[i].scatter(x_data, y_data, s=bubble_sizes, 
                                        c=colors[2], alpha=0.6, edgecolors='white', linewidth=0.5)
                
                axes[i].set_xlabel(f'{x_var}', fontproperties=self.font_prop)
                axes[i].set_ylabel(f'{y_var}', fontproperties=self.font_prop)
                axes[i].set_title(f'{name} - {x_var} vs {y_var}\n（气泡大小表示{size_var}）', 
                                fontproperties=self.font_prop)
                axes[i].grid(True, alpha=0.3)
        
        plt.tight_layout()
        plt.savefig(f'{self.output_dir}/散点气泡图_对比.png', bbox_inches='tight', dpi=300)
        plt.close()
        
        # 分别生成子图
        for i, (data, name, colors) in enumerate(datasets):
            if all(var in data.columns for var in [x_var, y_var, size_var]):
                fig, ax = plt.subplots(figsize=(8, 6))
                
                # 获取三列都非空的数据
                valid_data = data[[x_var, y_var, size_var]].dropna()
                if len(valid_data) == 0:
                    plt.close()
                    continue
                x_data = valid_data[x_var]
                y_data = valid_data[y_var]
                size_data = valid_data[size_var]
                
                size_normalized = (size_data - size_data.min()) / (size_data.max() - size_data.min())
                bubble_sizes = 20 + size_normalized * 200
                
                scatter = ax.scatter(x_data, y_data, s=bubble_sizes, 
                                   c=colors[2], alpha=0.6, edgecolors='white', linewidth=0.5)
                
                ax.set_xlabel(f'{x_var}', fontproperties=self.font_prop)
                ax.set_ylabel(f'{y_var}', fontproperties=self.font_prop)
                ax.set_title(f'{name} - {x_var} vs {y_var}\n（气泡大小表示{size_var}）', 
                           fontproperties=self.font_prop)
                ax.grid(True, alpha=0.3)
                
                plt.tight_layout()
                plt.savefig(f'{self.output_dir}/图{i+1}a_{name}_散点气泡图.png', bbox_inches='tight', dpi=300)
                plt.close()
        
        print("散点气泡图绘制完成！")
    
    def plot_diverging_bar(self):
        """绘制发散型条形图"""
        print("正在绘制发散型条形图...")
        self.set_plot_style()
        
        # 计算各变量的均值差异
        variables = ['年龄', '身高', '体重', '孕妇BMI', '检测孕周', 
                    '13号染色体的Z值', '18号染色体的Z值', '21号染色体的Z值']
        
        differences = []
        var_names = []
        
        for var in variables:
            if var in self.male_data.columns and var in self.female_data.columns:
                male_mean = self.male_data[var].mean()
                female_mean = self.female_data[var].mean()
                diff = male_mean - female_mean
                differences.append(diff)
                var_names.append(var)
        
        # 创建发散型条形图
        fig, ax = plt.subplots(figsize=(10, 8))
        
        y_pos = np.arange(len(var_names))
        colors = [self.colors_green[2] if x >= 0 else self.colors_blue[2] for x in differences]
        
        bars = ax.barh(y_pos, differences, color=colors, alpha=0.8, edgecolor='white', linewidth=0.5)
        
        # 添加零线
        ax.axvline(x=0, color='black', linestyle='-', linewidth=1)
        
        # 设置标签
        ax.set_yticks(y_pos)
        ax.set_yticklabels(var_names, fontproperties=self.font_prop)
        ax.set_xlabel('男胎均值 - 女胎均值', fontproperties=self.font_prop)
        ax.set_title('男胎与女胎各指标均值差异对比', fontproperties=self.font_prop, fontsize=14)
        
        # 添加数值标签
        for i, (bar, val) in enumerate(zip(bars, differences)):
            ax.text(val + (0.01 if val >= 0 else -0.01), i, f'{val:.3f}', 
                   ha='left' if val >= 0 else 'right', va='center', 
                   fontproperties=self.font_prop, fontsize=9)
        
        ax.grid(True, alpha=0.3, axis='x')
        
        plt.tight_layout()
        plt.savefig(f'{self.output_dir}/发散型条形图.png', bbox_inches='tight', dpi=300)
        plt.close()
        
        # 生成子图（按正负分组）
        positive_vars = [(var, diff) for var, diff in zip(var_names, differences) if diff >= 0]
        negative_vars = [(var, diff) for var, diff in zip(var_names, differences) if diff < 0]
        
        # 正值子图
        if positive_vars:
            fig, ax = plt.subplots(figsize=(8, 6))
            vars_pos, diffs_pos = zip(*positive_vars)
            y_pos = np.arange(len(vars_pos))
            
            bars = ax.barh(y_pos, diffs_pos, color=self.colors_green[2], alpha=0.8, 
                          edgecolor='white', linewidth=0.5)
            
            ax.set_yticks(y_pos)
            ax.set_yticklabels(vars_pos, fontproperties=self.font_prop)
            ax.set_xlabel('男胎均值 - 女胎均值', fontproperties=self.font_prop)
            ax.set_title('男胎均值高于女胎的指标', fontproperties=self.font_prop)
            
            for i, (bar, val) in enumerate(zip(bars, diffs_pos)):
                ax.text(val + 0.01, i, f'{val:.3f}', ha='left', va='center', 
                       fontproperties=self.font_prop, fontsize=9)
            
            ax.grid(True, alpha=0.3, axis='x')
            plt.tight_layout()
            plt.savefig(f'{self.output_dir}/图1a_发散型条形图_正值.png', bbox_inches='tight', dpi=300)
            plt.close()
        
        # 负值子图
        if negative_vars:
            fig, ax = plt.subplots(figsize=(8, 6))
            vars_neg, diffs_neg = zip(*negative_vars)
            y_pos = np.arange(len(vars_neg))
            
            bars = ax.barh(y_pos, diffs_neg, color=self.colors_blue[2], alpha=0.8, 
                          edgecolor='white', linewidth=0.5)
            
            ax.set_yticks(y_pos)
            ax.set_yticklabels(vars_neg, fontproperties=self.font_prop)
            ax.set_xlabel('男胎均值 - 女胎均值', fontproperties=self.font_prop)
            ax.set_title('女胎均值高于男胎的指标', fontproperties=self.font_prop)
            
            for i, (bar, val) in enumerate(zip(bars, diffs_neg)):
                ax.text(val - 0.01, i, f'{val:.3f}', ha='right', va='center', 
                       fontproperties=self.font_prop, fontsize=9)
            
            ax.grid(True, alpha=0.3, axis='x')
            plt.tight_layout()
            plt.savefig(f'{self.output_dir}/图1b_发散型条形图_负值.png', bbox_inches='tight', dpi=300)
            plt.close()
        
        print("发散型条形图绘制完成！")
    
    def plot_histogram_comparison(self):
        """绘制连续变量直方图（多组对比）"""
        print("正在绘制直方图对比...")
        self.set_plot_style()
        
        variables = ['年龄', '身高', '体重', '孕妇BMI', '检测孕周']
        
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        axes = axes.flatten()
        
        for i, var in enumerate(variables):
            if var in self.male_data.columns and var in self.female_data.columns:
                male_data = self.male_data[var].dropna()
                female_data = self.female_data[var].dropna()
                
                # 绘制重叠直方图
                axes[i].hist(male_data, bins=30, alpha=0.7, color=self.colors_green[2], 
                           label='男胎', density=True, edgecolor='white', linewidth=0.5)
                axes[i].hist(female_data, bins=30, alpha=0.7, color=self.colors_blue[2], 
                           label='女胎', density=True, edgecolor='white', linewidth=0.5)
                
                axes[i].set_xlabel(f'{var}', fontproperties=self.font_prop)
                axes[i].set_ylabel('密度', fontproperties=self.font_prop)
                axes[i].set_title(f'{var}分布对比', fontproperties=self.font_prop)
                axes[i].legend(prop=self.font_prop)
                axes[i].grid(True, alpha=0.3)
        
        # 隐藏多余子图
        axes[-1].set_visible(False)
        
        plt.tight_layout()
        plt.savefig(f'{self.output_dir}/直方图_多组对比.png', bbox_inches='tight', dpi=300)
        plt.close()
        
        # 分别生成子图
        for i, var in enumerate(variables):
            if var in self.male_data.columns and var in self.female_data.columns:
                fig, ax = plt.subplots(figsize=(8, 6))
                
                male_data = self.male_data[var].dropna()
                female_data = self.female_data[var].dropna()
                
                ax.hist(male_data, bins=30, alpha=0.7, color=self.colors_green[2], 
                       label='男胎', density=True, edgecolor='white', linewidth=0.5)
                ax.hist(female_data, bins=30, alpha=0.7, color=self.colors_blue[2], 
                       label='女胎', density=True, edgecolor='white', linewidth=0.5)
                
                ax.set_xlabel(f'{var}', fontproperties=self.font_prop)
                ax.set_ylabel('密度', fontproperties=self.font_prop)
                ax.set_title(f'{var}分布对比', fontproperties=self.font_prop)
                ax.legend(prop=self.font_prop)
                ax.grid(True, alpha=0.3)
                
                plt.tight_layout()
                plt.savefig(f'{self.output_dir}/图{i+1}a_{var}_直方图对比.png', bbox_inches='tight', dpi=300)
                plt.close()
        
        print("直方图对比绘制完成！")
    
    def plot_2d_contour(self):
        """绘制二维渐变色地形图"""
        print("正在绘制二维渐变色地形图...")
        self.set_plot_style()
        
        # 选择两个连续变量
        x_var = '年龄'
        y_var = '孕妇BMI'
        
        datasets = [(self.male_data, '男胎', self.cmap_green), 
                   (self.female_data, '女胎', self.cmap_blue)]
        
        fig, axes = plt.subplots(1, 2, figsize=(14, 6))
        
        for i, (data, name, cmap) in enumerate(datasets):
            if x_var in data.columns and y_var in data.columns:
                # 获取两列都非空的数据
                valid_data = data[[x_var, y_var]].dropna()
                if len(valid_data) == 0:
                    continue
                x_data = valid_data[x_var]
                y_data = valid_data[y_var]
                
                # 创建网格
                x_min, x_max = x_data.min(), x_data.max()
                y_min, y_max = y_data.min(), y_data.max()
                
                xx, yy = np.meshgrid(np.linspace(x_min, x_max, 50),
                                   np.linspace(y_min, y_max, 50))
                
                # 计算密度（使用高斯核密度估计）
                from scipy.stats import gaussian_kde
                
                positions = np.vstack([xx.ravel(), yy.ravel()])
                values = np.vstack([x_data, y_data])
                kernel = gaussian_kde(values)
                density = np.reshape(kernel(positions).T, xx.shape)
                
                # 绘制等高线图
                contour = axes[i].contourf(xx, yy, density, levels=20, cmap=cmap, alpha=0.8)
                axes[i].contour(xx, yy, density, levels=20, colors='white', linewidths=0.5, alpha=0.6)
                
                # 添加散点
                axes[i].scatter(x_data, y_data, c='white', s=1, alpha=0.3)
                
                axes[i].set_xlabel(f'{x_var}', fontproperties=self.font_prop)
                axes[i].set_ylabel(f'{y_var}', fontproperties=self.font_prop)
                axes[i].set_title(f'{name} - {x_var} vs {y_var} 密度分布', fontproperties=self.font_prop)
                
                # 添加颜色条
                cb = plt.colorbar(contour, ax=axes[i])
                cb.set_label('密度', fontproperties=self.font_prop)
        
        plt.tight_layout()
        plt.savefig(f'{self.output_dir}/二维渐变色地形图_对比.png', bbox_inches='tight', dpi=300)
        plt.close()
        
        # 分别生成子图
        for i, (data, name, cmap) in enumerate(datasets):
            if x_var in data.columns and y_var in data.columns:
                fig, ax = plt.subplots(figsize=(8, 6))
                
                # 获取两列都非空的数据
                valid_data = data[[x_var, y_var]].dropna()
                if len(valid_data) == 0:
                    plt.close()
                    continue
                x_data = valid_data[x_var]
                y_data = valid_data[y_var]
                
                x_min, x_max = x_data.min(), x_data.max()
                y_min, y_max = y_data.min(), y_data.max()
                
                xx, yy = np.meshgrid(np.linspace(x_min, x_max, 50),
                                   np.linspace(y_min, y_max, 50))
                
                from scipy.stats import gaussian_kde
                
                positions = np.vstack([xx.ravel(), yy.ravel()])
                values = np.vstack([x_data, y_data])
                kernel = gaussian_kde(values)
                density = np.reshape(kernel(positions).T, xx.shape)
                
                contour = ax.contourf(xx, yy, density, levels=20, cmap=cmap, alpha=0.8)
                ax.contour(xx, yy, density, levels=20, colors='white', linewidths=0.5, alpha=0.6)
                ax.scatter(x_data, y_data, c='white', s=1, alpha=0.3)
                
                ax.set_xlabel(f'{x_var}', fontproperties=self.font_prop)
                ax.set_ylabel(f'{y_var}', fontproperties=self.font_prop)
                ax.set_title(f'{name} - {x_var} vs {y_var} 密度分布', fontproperties=self.font_prop)
                
                cb = plt.colorbar(contour, ax=ax)
                cb.set_label('密度', fontproperties=self.font_prop)
                
                plt.tight_layout()
                plt.savefig(f'{self.output_dir}/图{i+1}a_{name}_二维渐变色地形图.png', bbox_inches='tight', dpi=300)
                plt.close()
        
        print("二维渐变色地形图绘制完成！")
    
    def generate_all_plots(self):
        """生成所有图表"""
        if not self.load_processed_data():
            print("数据加载失败，无法生成图表")
            return False
        
        print("开始生成所有图表...")
        
        try:
            self.plot_violin_comparison()
            self.plot_hexbin_subplots()
            self.plot_scatter_bubble()
            self.plot_diverging_bar()
            self.plot_histogram_comparison()
            self.plot_2d_contour()
            
            print(f"\n所有图表生成完成！保存在目录: {self.output_dir}")
            return True
            
        except Exception as e:
            print(f"图表生成过程中出现错误: {e}")
            return False

def main():
    """主函数"""
    print("开始科学绘图...")
    
    # 创建绘图器实例
    plotter = ScientificPlotter()
    
    # 生成所有图表
    if plotter.generate_all_plots():
        print("科学绘图完成！")
    else:
        print("科学绘图失败！")

if __name__ == "__main__":
    main()
